import os

from langchain_chroma import Chroma
from langchain_core.documents import Document
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnableLambda, RunnablePassthrough
from langchain_openai import OpenAIEmbeddings, AzureOpenAIEmbeddings, AzureChatOpenAI

os.environ["LANGCHAIN_TRACING_V2"] = "true"
os.environ["LANGCHAIN_API_KEY"] = "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"

os.environ["AZURE_OPENAI_ENDPOINT"] = "http://menshen.test.xdf.cn"
# os.environ["OPENAI_API_BASE"] = "http://menshen.test.xdf.cn"
os.environ["OPENAI_API_KEY"] = "c8575027653b42b1b47747f0b4ab135b"
os.environ["OPENAI_API_TYPE"] = "azure"
os.environ["OPENAI_API_VERSION"] = "2023-05-15"

llm = AzureChatOpenAI(
    deployment_name="gpt-4o",
    model_name="gpt-4o",
    temperature=0
)

documents = [
    Document(
        page_content="狗是伟大的伴侣，以其忠诚和友好而闻名",
        metadata={"source": "哺乳动物宠物文档"}
    ),
    Document(
        page_content="猫式独立的动物，通常喜欢自己的空间",
        metadata={"source": "哺乳动物宠物文档"}
    ),
    Document(
        page_content="金鱼是初学者的流行宠物，需要相对简单的护理",
        metadata={"source": "鱼类宠物文档"}
    ),
    Document(
        page_content="鹦鹉是聪明的鸟类，能够模仿人类的语言",
        metadata={"source": "鸟类宠物文档"}
    ),
    Document(
        page_content="兔子是社交动物，需要足够的空间跳跃",
        metadata={"source": "哺乳动物宠物文档"}
    )
]

# 实例化向量空间
vector_store = Chroma.from_documents(documents, embedding=AzureOpenAIEmbeddings())

# 分数越低，相似度越高
# print(vector_store.similarity_search_with_score("哪种宠物最适合作为宠物？", k=2))

retriever = RunnableLambda(vector_store.similarity_search).bind(k=1)

# print(retriever.batch(["家里养只猫怎么样？", "家里养条狗怎么样？"]))

message = """
    请结合上下文回答问题：
    {question}
    上下文：
    {context}
"""

prompt = ChatPromptTemplate.from_template(message)

chain = {"question": RunnablePassthrough(), "context": retriever} | prompt | llm

resp = chain.invoke("请举例说明猪的智商高")

print(resp.content)
